#!/usr/bin/env python

'''
ec2-stop-cluster.py - Script to shutdown MPI cluster on EC2, change the configuration strings in 
EC2config.py to match your AWS keys

Created by Peter Skomoroch on 2007-04-09.
Copyright (c) 2007 DataWrangling. All rights reserved.

'''

import sys
import EC2
import time
from EC2config import *
import getopt


help_message = '''
ec2-stop-cluster.py - Script to shutdown MPI cluster on EC2, change the configuration strings in 
EC2config.py to match your AWS keys
'''


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


def main(argv=None):
    if argv is None:
        argv = sys.argv
    try:
        try:
            opts, args = getopt.getopt(argv[1:], "ho:v", ["help", "output="])
        except getopt.error, msg:
            raise Usage(msg)
    
        # option processing
        for option, value in opts:
            if option == "-v":
                verbose = True
            if option in ("-h", "--help"):
                raise Usage(help_message)
            if option in ("-o", "--output"):
                output = value
                
        stopcluster()       
    
    except Usage, err:
        print >> sys.stderr, sys.argv[0].split("/")[-1] + ": " + str(err.msg)
        print >> sys.stderr, "\t for help use --help"
        return 2

    
def stopcluster():  
    resp = raw_input("This will stop all your ElasticWulf Instances, are you sure (yes/no)? ")
    if resp == 'yes':
        conn = EC2.AWSAuthConnection(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)
        if IMAGE_ID:
            print "worker image",IMAGE_ID
        if globals().has_key("MASTER_IMAGE_ID"):
            print "master image",MASTER_IMAGE_ID
        instance_response=conn.describe_instances()
        print "----- listing instances -----"
        print instance_response
        parsed_response=instance_response.parse()  

        running_instances=[]
        for chunk in parsed_response:
            if chunk[0]=='INSTANCE':
                if chunk[-1]=='running':
                    if (chunk[2] == IMAGE_ID or chunk[2] == MASTER_IMAGE_ID):
                        running_instances.append(chunk[1])
        if len(running_instances) > 0:
            print "\n---- Stopping instance Id's ----"
            for instance in running_instances:
                print "Stoping Instance Id = %s " % instance
            print "\nWaiting for shutdown ...."
            conn.terminate_instances(running_instances)
            time.sleep(5)

        instance_response=conn.describe_instances()
        print "----- listing new state of instances -----"
        print instance_response
    
    else:
        print "exiting without stopping instances...."

if __name__ == "__main__":
    sys.exit(main())

